/*
 * Description: wait user logon in windows
 */
#include "session_utils.h"
#include <WtsApi32.h>
#include <tchar.h>
#include "log.h"
#include "guest_utils.h"
#include "explorer_style/explorer_style_modifier.h"

#pragma comment(lib, "iphlpapi.lib")

namespace COMMON {
std::vector<SessionUtils::SessionCallback> SessionUtils::m_logonCallbacks = {};
std::map<std::string, SessionUtils::SessionChangeCallback> SessionUtils::m_sessionCallbackMap_ = {};
std::mutex SessionUtils::m_mutex_ = {};
bool SessionUtils::m_flag = false;

void TellBalloon()
{
    LOG_INFO("Start notify balloon");
    HKEY hKey = NULL;
    // 创建注册项
    long result = RegCreateKeyEx(HKEY_LOCAL_MACHINE, _T("SOFTWARE\\DesktopAppEngine"), 0, NULL,
        REG_OPTION_VOLATILE, KEY_ALL_ACCESS, NULL, &hKey, NULL);
    if (result != ERROR_SUCCESS) {
        LOG_ERR("Reg add failed: {}", result);
        return;
    }
    // 写注册表
    const DWORD newValue = 1;
    DWORD size = sizeof(newValue);
    result = RegSetValueEx(hKey, _T("IsReady"), 0, REG_DWORD, LPBYTE(&newValue), size);
    if (result != ERROR_SUCCESS) {
        LOG_ERR("Reg set failed: {}", result);
        RegCloseKey(hKey);
        return;
    }
    LOG_INFO("Notify balloon success");
    // 关闭注册表
    RegCloseKey(hKey);
    return;
}

void SessionUtils::Init()
{
    if (IsUserLogon()) {
        TellBalloon();
        EXPLORER_STYLE_MODIFIER.CheckCurrentUserExplorerStyle();
        m_flag = true;
    }
    std::thread t1(SessionUtils::SessionChange);
    t1.detach();
    RegLogonCallback(std::bind(
        &GUEST::ExplorerStyleModifier::CheckCurrentUserExplorerStyle, &EXPLORER_STYLE_MODIFIER));
    RegLogonCallback(SetQuickStartStatus);
    LOG_INFO("Init success");
}

void SessionUtils::SetQuickStartStatus()
{
    HANDLE hServer = WTSOpenServer(NULL);
    if (hServer == NULL) {
        LOG_ERR("Open server failed");
        return;
    }
    PWTS_SESSION_INFO pSessionInfo = nullptr;
    DWORD count = 0;
    if (!WTSEnumerateSessions(hServer, 0, 1, &pSessionInfo, &count)) {
        LOG_ERR("Enumerate sessions failed");
        WTSCloseServer(hServer);
        return;
    }
    int userNum = 0;
    for (DWORD i = 0; i < count; i++) {
        WTS_INFO_CLASS infoClass = WTSUserName;
        LPSTR userName = NULL;
        DWORD bytesReturned = 0;
        if (!WTSQuerySessionInformation(
                WTS_CURRENT_SERVER_HANDLE, pSessionInfo[i].SessionId, infoClass, &userName, &bytesReturned)) {
            LOG_ERR("GetInfo failed:{}", GetLastError());
            continue;
        }
        if (userName == NULL) {
            LOG_ERR("Get user name failed");
            continue;
        }
        if (strcmp(userName, "") != 0) {
            userNum++;
        } else {
            LOG_INFO("Session is belong to system");
        }
        WTSFreeMemory(userName);
    }
    if (userNum > 1) {
        std::string disableQuickStartCmd = "powercfg /h off";
        if (!GuestUtils::ExecCmd(disableQuickStartCmd, false)) {
            LOG_ERR("Failed to disable quick start");
        }
    } else {
        std::string enableQuickStartCmd = "powercfg /h on";
        if (!GuestUtils::ExecCmd(enableQuickStartCmd, false)) {
            LOG_ERR("Failed to enable quick start");
        }
    }
    WTSFreeMemory(pSessionInfo);
    WTSCloseServer(hServer);
}

void SessionUtils::Lock()
{
    LOG_INFO("Lock session");
    std::string lockCmd = "rundll32.exe user32.dll,LockWorkStation";
    if (IsUserLogon()) {
        if (!GuestUtils::ExecCmd(lockCmd, true)) {
            LOG_ERR("Failed to execute {}", lockCmd);
        }
    }
}

void SessionUtils::dealMsg(MSG msg)
{
    LOG_INFO("Session event: {}", msg.wParam);
    if (!SessionUtils::m_sessionCallbackMap_.empty()) {
        int32_t state = static_cast<int32_t>(msg.wParam);
        std::unique_lock<std::mutex> lock(m_mutex_);
        for (auto it = SessionUtils::m_sessionCallbackMap_.begin();
            it != SessionUtils::m_sessionCallbackMap_.end(); it++) {
            it->second(state);
        }
        lock.unlock();
    }
    TranslateMessage(&msg);
    DispatchMessage(&msg);
}

void SessionUtils::SessionChange()
{
    LOG_INFO("START WATCH SESSION");
    MSG msg;
    BOOL ret;
    HWND hWnd = CreateWindow("STATIC", "", WS_POPUP, 0, 0, 0, 0, NULL, NULL, GetModuleHandle(NULL), NULL);
    if (!hWnd) {
        LOG_ERR("CreateWindow failed: {}", GetLastError());
        return;
    }
    ShowWindow(hWnd, SW_HIDE);
    if (!WTSRegisterSessionNotification(hWnd, NOTIFY_FOR_ALL_SESSIONS)) {
        LOG_ERR("Failed to register window notify: {}", GetLastError());
        if (!DestroyWindow(hWnd)) {
            LOG_ERR("Destroy window failed: {}", GetLastError());
        }
        return;
    }
    // 订阅发送到本窗口的SessionChange消息
    while ((ret = GetMessage(&msg, hWnd, WM_WTSSESSION_CHANGE, WM_WTSSESSION_CHANGE)) != 0) {
        if (ret == -1) {
            LOG_ERR("Get msg failed: {}", GetLastError());
            break;
        }
        LOG_INFO("sessionchange msg: {}", msg.message);
        if (msg.message != WM_WTSSESSION_CHANGE) {
            continue;
        }
        if (msg.wParam == WTS_SESSION_LOGON) {
            for (auto fn: m_logonCallbacks) {
                fn();
            }
            if (!m_flag) {
                // 本次启动中首次登陆通知Balloon
                TellBalloon();
                m_flag = true;
            }
        }
        if (msg.wParam == WTS_CONSOLE_CONNECT) {
            SetQuickStartStatus();
        }
        dealMsg(msg);
    }
    LOG_INFO("Message end");
    WTSUnRegisterSessionNotification(hWnd);
    if (!DestroyWindow(hWnd)) {
        LOG_ERR("Destroy window failed: {}", GetLastError());
    }
}

bool SessionUtils::IsUserLogon()
{
    DWORD sessionId = WTSGetActiveConsoleSessionId();
    LOG_INFO("Get Session {}", sessionId);
    if (sessionId == 0xFFFFFFFF) {
        LOG_INFO("Nobody logon");
        return false;
    } else {
        WTS_INFO_CLASS infoClass = WTSUserName;
        LPSTR userName = NULL;
        DWORD bytesReturned = 0;
        if (WTSQuerySessionInformation(WTS_CURRENT_SERVER_HANDLE, sessionId, infoClass, &userName, &bytesReturned)) {
            if (userName == NULL) {
                LOG_ERR("Get user name failed");
                return false;
            }
            if (strcmp(userName, "") != 0) {
                LOG_INFO("User {} has logon", userName);
                WTSFreeMemory(userName);
                return true;
            } else {
                LOG_INFO("Session is belong to system");
                WTSFreeMemory(userName);
                return false;
            }
        } else {
            LOG_ERR("GetInfo failed:{}", GetLastError());
        }
    }
    return true;
}

void SessionUtils::RegSessionChangeCallback(const std::string& name, const SessionChangeCallback& fn)
{
    if (name.empty()) {
        return;
    }
    std::unique_lock<std::mutex> lock(m_mutex_);
    m_sessionCallbackMap_.emplace(name, fn);
    LOG_INFO("{} register session change callback success", name);
}

void SessionUtils::UnRegSessionChangeCallback(const std::string& name)
{
    if (name.empty()) {
        return;
    }
    auto it = m_sessionCallbackMap_.find(name);
    if (it != m_sessionCallbackMap_.end()) {
        std::unique_lock<std::mutex> lock(m_mutex_);
        m_sessionCallbackMap_.erase(name);
        LOG_INFO("{} unregister session change callback success", name);
    }
}
}